package security.keyManagement;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.security.KeyPair;
import java.security.KeyStore;
import java.security.KeyStore.PasswordProtection;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Security;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

public class KeyStoreManager {

	private KeyStore _keyStore;

	private String _keyStoreAddress;

	private PasswordProtection _passwordProtection;

	private static final String _originalCertcopKeyStoreAddress = "config/CERTCOP.keystore";
	
	private static final String _originalVericopKeyStoreAddress = "config/VERICOP.keystore";

	private static final String _originalClientKeyStoreAddress = "config/Client.keystore";

	private KeyStoreManager(boolean server, boolean certcop){
		if(server){
			if(certcop)
				_keyStoreAddress = _originalCertcopKeyStoreAddress;
			else
				_keyStoreAddress = _originalVericopKeyStoreAddress;
		}			
		else
			_keyStoreAddress = _originalClientKeyStoreAddress;
	}

	public KeyStoreManager(String password, boolean server, boolean certcop) throws IOException{
		if(server){
			if(certcop)
				_keyStoreAddress = _originalCertcopKeyStoreAddress;
			else
				_keyStoreAddress = _originalVericopKeyStoreAddress;
		}else
			_keyStoreAddress = _originalClientKeyStoreAddress;

		try {
			FileInputStream fis = new FileInputStream(_keyStoreAddress);
			_keyStore = KeyStore.getInstance("JCEKS");

			if(password != null){
				_keyStore.load(fis, password.toCharArray());
				_passwordProtection = new PasswordProtection(password.toCharArray());
			}
			else
				_keyStore.load(fis, "".toCharArray());			

		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} 
	}
	
	public String getPassword(){
		return String.valueOf(_passwordProtection.getPassword());
	}
	
	public PrivateKey getServerPrivateKey(boolean certcop) {
		if(certcop){
			try {
				return (PrivateKey) _keyStore.getKey("CERTCOPServerKeyPair", _passwordProtection.getPassword());
			} catch (UnrecoverableKeyException e) {
				e.printStackTrace();
			} catch (KeyStoreException e) {
				e.printStackTrace();
			} catch (NoSuchAlgorithmException e) {
				e.printStackTrace();
			}
		}else{
			try {
				return (PrivateKey) _keyStore.getKey("VERICOPServerKeyPair", _passwordProtection.getPassword());
			} catch (UnrecoverableKeyException e) {
				e.printStackTrace();
			} catch (KeyStoreException e) {
				e.printStackTrace();
			} catch (NoSuchAlgorithmException e) {
				e.printStackTrace();
			}
		}
		return null;
	}
	
	public KeyPair getVericopKeyPair(){
		try {
			PrivateKey privateKey =  (PrivateKey) _keyStore.getKey("VERICOPServerKeyPair", _passwordProtection.getPassword());
			Certificate vericopCert = _keyStore.getCertificateChain("VERICOPServerKeyPair")[0];
			
			KeyPair keyPair = new KeyPair(vericopCert.getPublicKey(), privateKey);
			return keyPair;
		} catch (UnrecoverableKeyException e) {
			e.printStackTrace();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		}
		
		return null;
	}
	
	public PublicKey getVericopPublicKey() {
		try {
			Certificate vericopCert = _keyStore.getCertificate("VERICOPServerCert");			
			return vericopCert.getPublicKey();
		} catch (KeyStoreException e) {
			e.printStackTrace();
		}
		
		return null;
	}
	
	public PublicKey getCertcopPublicKey(){
		try {
			Certificate certcopCert = _keyStore.getCertificate("CERTCOPServerCert");
			return certcopCert.getPublicKey();
		} catch (KeyStoreException e) {
			e.printStackTrace();
			return null;
		}
	}
	
	public PublicKey getWebServerPublicKey() {
		try {
			Certificate webServerCert = _keyStore.getCertificate("WebServerCert");
			return webServerCert.getPublicKey();
		} catch (KeyStoreException e) {
			e.printStackTrace();
			return null;
		}
	}
	
	// TODO Verificar se correcto
	public String getCertcopCertSerial() {
	    try {
            X509Certificate certcopCert = (X509Certificate) _keyStore.getCertificateChain("CERTCOPServerKeyPair")[0];
            return certcopCert.getSerialNumber().toString();
        } catch (KeyStoreException e) {
            e.printStackTrace();
            return null;
        }
	}
	
	protected void finalize() throws Throwable{
		try{
			FileOutputStream fos = new FileOutputStream(_keyStoreAddress);

			_keyStore.store(fos,_passwordProtection.getPassword());

			fos.flush();
			fos.close();

			_passwordProtection.destroy();			
		}finally{
			super.finalize();
		}
	}

	/** This method should never be called. Only to keyStore generation
	 * purposes.
	 * @throws KeyStoreException 
	 * @throws IOException 
	 * @throws CertificateException 
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeySpecException 
	 */
	private void createCertcopKeyStore() throws KeyStoreException, NoSuchAlgorithmException, CertificateException, IOException, InvalidKeySpecException{
		File file = new File(_keyStoreAddress);
		file.delete();

		Security.addProvider(new BouncyCastleProvider());		

		Certificate rootCaCert = Utils.getX509CertFromFile("resources/ca-cert.pem");

		Certificate certcopCert = Utils.getX509CertFromFile("resources/certcopserver-cert.pem");
		PrivateKey privateKey = Utils.getPrivateKeyFromFile("resources/certcopserver-key.pem");

		Certificate dbServerCert = Utils.getX509CertFromFile("resources/dbserver-cert.pem");
		
		Certificate vericopServerCert = Utils.getX509CertFromFile("resources/vericopserver-cert.pem");
		
		Certificate webServerCert = Utils.getX509CertFromFile("resources/webserver-cert.pem");

		PBEKeySpec keySpec = new PBEKeySpec("P3TRdcXkS5".toCharArray());
		SecretKeyFactory skf = SecretKeyFactory.getInstance("PBEWithMD5AndDES");
		SecretKey mySqlKey = skf.generateSecret(keySpec);

		_keyStore = KeyStore.getInstance("JCEKS");
		_keyStore.load(null, null);

		Certificate[] certcopChain = new Certificate[]{certcopCert, rootCaCert};
		_keyStore.setKeyEntry("CERTCOPServerKeyPair", privateKey, "W2g4xmF8kc".toCharArray(), certcopChain);

		_keyStore.setCertificateEntry("DBServerCert", dbServerCert);
		
		_keyStore.setCertificateEntry("VERICOPServerCert", vericopServerCert);
		
		_keyStore.setCertificateEntry("WebServerCert", webServerCert);

		_keyStore.setKeyEntry("MySQLPassword", mySqlKey, "W2g4xmF8kc".toCharArray(), null);		

		OutputStream fos = new FileOutputStream(_originalCertcopKeyStoreAddress);
		_keyStore.store(fos,"W2g4xmF8kc".toCharArray());

		fos.flush();
		fos.close();		
	}
	
	/** This method should never be called. Only to keyStore generation
	 * purposes.
	 * @throws KeyStoreException 
	 * @throws IOException 
	 * @throws CertificateException 
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeySpecException 
	 */
	private void createVericopKeyStore() throws KeyStoreException, NoSuchAlgorithmException, CertificateException, IOException, InvalidKeySpecException{
		File file = new File(_keyStoreAddress);
		file.delete();

		Security.addProvider(new BouncyCastleProvider());		

		Certificate rootCaCert = Utils.getX509CertFromFile("resources/ca-cert.pem");

		Certificate vericopCert = Utils.getX509CertFromFile("resources/vericopserver-cert.pem");
		PrivateKey privateKey = Utils.getPrivateKeyFromFile("resources/vericopserver-key.pem");

		Certificate dbServerCert = Utils.getX509CertFromFile("resources/dbserver-cert.pem");
		
		Certificate certcopServerCert = Utils.getX509CertFromFile("resources/certcopserver-cert.pem");

		PBEKeySpec keySpec = new PBEKeySpec("P3TRdcXkS5".toCharArray());
		SecretKeyFactory skf = SecretKeyFactory.getInstance("PBEWithMD5AndDES");
		SecretKey mySqlKey = skf.generateSecret(keySpec);

		_keyStore = KeyStore.getInstance("JCEKS");
		_keyStore.load(null, null);

		Certificate[] vericopChain = new Certificate[]{vericopCert, rootCaCert};
		_keyStore.setKeyEntry("VERICOPServerKeyPair", privateKey, "Jq8C531GCt".toCharArray(), vericopChain);

		_keyStore.setCertificateEntry("DBServerCert", dbServerCert);
		
		_keyStore.setCertificateEntry("CERTCOPServerCert", certcopServerCert);

		_keyStore.setKeyEntry("MySQLPassword", mySqlKey, "Jq8C531GCt".toCharArray(), null);		

		OutputStream fos = new FileOutputStream(_originalVericopKeyStoreAddress);
		_keyStore.store(fos,"Jq8C531GCt".toCharArray());

		fos.flush();
		fos.close();		
	}

	/** This method should never be called. Only to keyStore generation
	 * purposes.
	 * @throws KeyStoreException 
	 * @throws IOException 
	 * @throws CertificateException 
	 * @throws NoSuchAlgorithmException 
	 */
	private void createClientKeyStore() throws KeyStoreException, NoSuchAlgorithmException, CertificateException, IOException{
		File file = new File(_keyStoreAddress);
		file.delete();

		Security.addProvider(new BouncyCastleProvider());

		Certificate rootCaCert = Utils.getX509CertFromFile("resources/ca-cert.pem");

		Certificate certcopCert = Utils.getX509CertFromFile("resources/certcopserver-cert.pem");
		
		Certificate vericopServerCert = Utils.getX509CertFromFile("resources/vericopserver-cert.pem");

		_keyStore = KeyStore.getInstance("JCEKS");
		_keyStore.load(null, null);

		_keyStore.setCertificateEntry("RootCACert", rootCaCert);

		_keyStore.setCertificateEntry("CERTCOPServerCert", certcopCert);
		
		_keyStore.setCertificateEntry("VERICOPServerCert", vericopServerCert);

		OutputStream fos = new FileOutputStream(_originalClientKeyStoreAddress);
		_keyStore.store(fos,"".toCharArray());

		fos.flush();
		fos.close();

	}
	
	/**
	 * DO NOT USE. ONLY FOR KEYSTORE GENERATION PURPOSES!!
	 * @param args
	 * @throws Throwable
	 */
	public static void main(String[] args) throws Throwable{
		
//		KeyStoreManager serverManager = new KeyStoreManager(true, true);
//		serverManager.createCertcopKeyStore();
//		
//		System.out.println(Utils.getReadableX509Certificate(serverManager._keyStore.getCertificate("DBServerCert")).getSerialNumber());
//		
//		System.out.println(Utils.getReadableX509Certificate(serverManager._keyStore.getCertificate("VERICOPServerCert")).getSerialNumber());
//		
//		System.out.println(Utils.getReadableX509Certificate(serverManager._keyStore.getCertificate("WebServerCert")).getSerialNumber());
//
//		KeyStoreManager manager = new KeyStoreManager(false, false);
//		manager.createClientKeyStore();
//		
//		KeyStoreManager vericopManager = new KeyStoreManager(true, false);
//		vericopManager.createVericopKeyStore();
	    
	    // Access keystore certs
	    KeyStoreManager keyStoreManager = new KeyStoreManager("W2g4xmF8kc", true, true);
	    
	    System.out.println(keyStoreManager.getCertcopCertSerial());
	}
}
